import xgboost as xgb
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np

"""
XGBoost 多分类示例 - 使用鸢尾花数据集
鸢尾花数据集包含3个类别：0-setosa, 1-versicolor, 2-virginica
特征数：4个（花萼长度、花萼宽度、花瓣长度、花瓣宽度）
样本数：150个
"""

# 加载鸢尾花数据集
X, y = load_iris(return_X_y=True)

# 展示数据类型和基本信息
print("X的数据类型:", type(X))
print("y的数据类型:", type(y))
print("X的形状:", X.shape)
print("y的形状:", y.shape)

# 展示数据集信息
print("数据集特征数:", X.shape[1])
print("数据集样本数:", X.shape[0])
print("类别分布:", np.bincount(y))
print("类别标签:", np.unique(y))
print("y的前五个值:", y[:5])

# 划分数据集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42, stratify=y
)


def train_model(X_train, y_train, max_depth=3, n_estimators=100):
    """
    训练XGBoost多分类模型
    :param X_train: 训练特征数据
    :param y_train: 训练目标数据
    :param max_depth: 树的最大深度
    :param n_estimators: 树的数量
    :return: 训练好的模型实例
    """
    # 使用multi:softprob目标函数进行多分类
    clf = xgb.XGBClassifier(
        objective="multi:softprob",
        max_depth=max_depth,
        n_estimators=n_estimators,
        num_class=3,  # 指定类别数量
        random_state=42,
    )
    clf.fit(X_train, y_train)
    return clf


def evaluate_model(model, X_test, y_test):
    """
    评估模型性能
    :param model: 训练好的模型实例
    :param X_test: 测试集特征
    :param y_test: 测试集真实标签
    :return: 准确率
    """
    # 预测
    y_pred = model.predict(X_test)
    # y_pred_proba = model.predict_proba(X_test)

    # 计算准确率
    accuracy = accuracy_score(y_test, y_pred)
    print(f"准确率: {accuracy:.4f}")

    return accuracy


def save_model(model, model_path):
    """
    保存模型到指定路径
    :param model: 训练好的模型实例
    :param model_path: 模型保存路径
    """
    model.save_model(model_path)
    print(f"模型已保存至 {model_path}")


def load_model(model_path):
    """
    从指定路径加载模型
    :param model_path: 模型保存路径
    :return: 加载的模型实例
    """
    model = xgb.XGBClassifier()
    model.load_model(model_path)
    print(f"模型已从 {model_path} 加载")
    return model


def predict(model, X):
    """
    使用模型进行预测
    :param model: 训练好的模型实例
    :param X: 输入特征数据
    :return: 预测结果
    """
    return model.predict(X)


# 主程序执行
if __name__ == "__main__":
    # 训练模型
    print("开始训练模型...")
    model = train_model(X_train, y_train)

    # 评估模型
    print("模型评估结果:")
    accuracy = evaluate_model(model, X_test, y_test)

    # 保存模型
    # save_model(model, "iris_xgboost_classifier.json")

    # 加载模型
    # print("\n验证模型加载功能:")
    # loaded_model = load_model("diabetes_xgboost_model.json")

    # 使用模型，预测值是多少

    print("程序执行完成！")
